{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fd822472",
   "metadata": {},
   "source": [
    "本baseline使用的是yoloV5基本框架，推荐使用GPU进行训练，如果没有GPU加速，可以尝试去kaggle上面的kernel使用免费GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4154097",
   "metadata": {},
   "source": [
    "# 数据预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35bbe7ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "根据yoloV5的要求，将xml文件，转化为txt文件，同时将具体坐标转换为对应的比例。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99aba4aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xml.dom.minidom import parse\n",
    "import xml.dom.minidom\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9521a2f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "l=['knife',\n",
    " 'scissors',\n",
    " 'sharpTools',\n",
    " 'expandableBaton',\n",
    " 'smallGlassBottle',\n",
    " 'electricBaton',\n",
    " 'plasticBeverageBottle',\n",
    " 'plasticBottleWithaNozzle',\n",
    " 'electronicEquipment',\n",
    " 'battery',\n",
    " 'seal',\n",
    " 'umbrella']\n",
    "dic = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3884c209",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n,i in enumerate(l):\n",
    "    dic[i] = str(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0849ccb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \".\\\\初赛第一阶段_训练集\\\\domain1\\\\XML\\\\\"\n",
    "path2 = \".\\\\初赛第一阶段_训练集\\\\domain1\\\\label\\\\\"\n",
    "# 写入path2对应的文件夹"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61604073",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in os.listdir(path):\n",
    "    l0 = []\n",
    "    DOMTree = xml.dom.minidom.parse(path+i)\n",
    "    root = DOMTree.documentElement\n",
    "    \n",
    "    width = int(root.getElementsByTagName('size')[0].getElementsByTagName('width')[0].firstChild.data)\n",
    "    height = int(root.getElementsByTagName('size')[0].getElementsByTagName('height')[0].firstChild.data)\n",
    "    for j in root.getElementsByTagName('object'):\n",
    "        l1 = []\n",
    "        l1.append(dic[j.getElementsByTagName('name')[0].firstChild.data])\n",
    "        xmin = int(j.getElementsByTagName('bndbox')[0].getElementsByTagName('xmin')[0].firstChild.data)\n",
    "        ymin = int(j.getElementsByTagName('bndbox')[0].getElementsByTagName('ymin')[0].firstChild.data)\n",
    "        xmax = int(j.getElementsByTagName('bndbox')[0].getElementsByTagName('xmax')[0].firstChild.data)\n",
    "        ymax = int(j.getElementsByTagName('bndbox')[0].getElementsByTagName('ymax')[0].firstChild.data)\n",
    "        l1.append(str((xmin+xmax)/2/width))\n",
    "        l1.append(str((ymin+ymax)/2/height))\n",
    "        l1.append(str((xmax-xmin)/width))\n",
    "        l1.append(str((ymax-ymin)/height))\n",
    "        l0.append(l1)\n",
    "    with open(path2 + i.split('.')[0]+'.txt','w',encoding='utf-8') as f:\n",
    "        [f.write('{0}\\n'.format(' '.join(value))) for value in l0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6029cea6",
   "metadata": {},
   "source": [
    "# 训练"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d52232ea",
   "metadata": {},
   "source": [
    "## kaggle或linux版本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96858825",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This Python 3 environment comes with many helpful analytics libraries installed\n",
    "# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n",
    "# For example, here's several helpful packages to load\n",
    "\n",
    "import numpy as np # linear algebra\n",
    "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "\n",
    "# Input data files are available in the read-only \"../input/\" directory\n",
    "# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n",
    "\n",
    "import os\n",
    "\n",
    "# for dirname, _, filenames in os.walk('/kaggle/input'):\n",
    "#     for filename in filenames:\n",
    "#         print(os.path.join(dirname, filename))\n",
    "\n",
    "# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n",
    "# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "anonymous-politics",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-07T12:06:11.669519Z",
     "iopub.status.busy": "2021-07-07T12:06:11.669029Z",
     "iopub.status.idle": "2021-07-07T12:06:23.093844Z",
     "shell.execute_reply": "2021-07-07T12:06:23.094424Z",
     "shell.execute_reply.started": "2021-07-07T07:21:30.752305Z"
    },
    "papermill": {
     "duration": 11.432947,
     "end_time": "2021-07-07T12:06:23.094582",
     "exception": false,
     "start_time": "2021-07-07T12:06:11.661635",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setup complete. Using torch 1.7.0 (CPU)\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/ultralytics/yolov5  # clone repo\n",
    "%cd yolov5\n",
    "%pip install -qr requirements.txt  # install dependencies\n",
    "\n",
    "import torch\n",
    "from IPython.display import Image, clear_output  # to display images\n",
    "\n",
    "clear_output()\n",
    "print(f\"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "protecting-challenge",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-15T03:20:33.437409Z",
     "start_time": "2021-07-15T03:20:33.424338Z"
    },
    "execution": {
     "iopub.execute_input": "2021-07-07T12:06:23.114793Z",
     "iopub.status.busy": "2021-07-07T12:06:23.112036Z",
     "iopub.status.idle": "2021-07-07T12:06:31.219167Z",
     "shell.execute_reply": "2021-07-07T12:06:31.217143Z",
     "shell.execute_reply.started": "2021-07-07T07:21:38.863407Z"
    },
    "papermill": {
     "duration": 8.118314,
     "end_time": "2021-07-07T12:06:31.219464",
     "exception": true,
     "start_time": "2021-07-07T12:06:23.101150",
     "status": "failed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Weights & Biases  (optional)\n",
    "%pip install -q wandb\n",
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "anonymous-bulgaria",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-07T07:22:04.388087Z",
     "iopub.status.busy": "2021-07-07T07:22:04.387754Z",
     "iopub.status.idle": "2021-07-07T12:04:05.571025Z",
     "shell.execute_reply": "2021-07-07T12:04:05.570042Z",
     "shell.execute_reply.started": "2021-07-07T07:22:04.388055Z"
    },
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python train.py --img 640 --batch 32 --epochs 100 --data /kaggle/input/findxray/kaggle.yaml --weights /kaggle/input/findxray/yolov5s.pt --nosave --cache\n",
    "# 其中，kaggle.yaml和yolov5s.pt文件的目录需要根据自己情况修改。最终结果输出到runs/train/exp文件夹下"
   ]
  },
  {
   "cell_type": "raw",
   "id": "69e41cb1",
   "metadata": {},
   "source": [
    "##########################################################\n",
    "# 附kaggle.yaml文件的内容,创建txt文件，将下面代码复制，并给文件重命名\n",
    "# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]\n",
    "train: /kaggle/input/findxray/xray/images/train2017/\n",
    "# 训练集图像路径，同时需要在对应的/kaggle/input/findxray/xray/labels/train2017/下放置labels，但不需要声明，会自动识别\n",
    "val: /kaggle/input/findxray/xray/images/val2017/\n",
    "# 检验集图像路径\n",
    "\n",
    "# number of classes\n",
    "nc: 12\n",
    "\n",
    "# class names\n",
    "names: ['knife',\n",
    " 'scissors',\n",
    " 'sharpTools',\n",
    " 'expandableBaton',\n",
    " 'smallGlassBottle',\n",
    " 'electricBaton',\n",
    " 'plasticBeverageBottle',\n",
    " 'plasticBottleWithaNozzle',\n",
    " 'electronicEquipment',\n",
    " 'battery',\n",
    " 'seal',\n",
    " 'umbrella']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rubber-victor",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python detect.py --weights ./runs/train/exp/weights/last.pt --img 640 --conf 0.5 --source ../test/ --save-txt  --save-conf\n",
    "# 其中，last.pt和test文件夹的目录需要根据自己情况修改"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c99d2ec5",
   "metadata": {},
   "source": [
    "## windows版本"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6efead07",
   "metadata": {},
   "source": [
    "请参考\n",
    "https://github.com/wudashuo/yolov5\n",
    "https://blog.csdn.net/qq_45701791/article/details/113992622\n",
    "下载依赖包，下载yoloV5代码，之后在powershell（命令行）当中运行训练及测试命令即可。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1ae7ef6",
   "metadata": {},
   "source": [
    "# 输出结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d98f8c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pandas as pd\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12036d09",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取图片大小\n",
    "df_size = pd.DataFrame(columns=['name','width','height'])\n",
    "\n",
    "path = 'E:\\\\xray\\\\test\\\\'\n",
    "for file in os.listdir(path):\n",
    "    img = cv2.imread(path+file)\n",
    "    size = img.shape\n",
    "    df_size.loc[df_size.shape[0]] = [file.split('.')[0],size[1],size[0]]\n",
    "    \n",
    "df_size.to_csv('size.csv',index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9059e5d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_size = pd.read_csv('size.csv')\n",
    "df_size =  df_size.set_index('name')\n",
    "df_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1040828c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 因为yoloV5最终预测的是比例，需要根据图像大小把比例还原成坐标\n",
    "\n",
    "result=[]\n",
    "path = 'E:\\\\xray\\\\yolov5-master\\\\runs\\\\detect\\\\exp\\\\labels\\\\'\n",
    "path2 = 'E:\\\\xray\\\\test\\\\'\n",
    "for file in os.listdir(path2):\n",
    "    pic = []\n",
    "    try:\n",
    "        df = pd.read_csv(path+file.split('.')[0]+'.txt',header=None,delimiter=' ')\n",
    "        df[0] = df[0]+1\n",
    "        df = df.set_index(0)\n",
    "\n",
    "        df[1] = df[1]*df_size.loc[file.split('.')[0],'width']\n",
    "        df[3] = df[3]*df_size.loc[file.split('.')[0],'width']\n",
    "\n",
    "        df[2] = df[2]*df_size.loc[file.split('.')[0],'height']\n",
    "        df[4] = df[4]*df_size.loc[file.split('.')[0],'height']\n",
    "        \n",
    "        df[1],df[2],df[3],df[4] = df[1]-df[3]/2,df[2]-df[4]/2,df[1]+df[3]/2,df[2]+df[4]/2\n",
    "        for i in range(1,13):\n",
    "            try:\n",
    "                pic.append(df.loc[i].values.tolist())\n",
    "            except:\n",
    "                pic.append([])\n",
    "    except:\n",
    "        for i in range(1,13):\n",
    "            pic.append([])\n",
    "    result.append(pic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bf39959",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输出文件\n",
    "with open('sub.json','w') as f:\n",
    "    f.write(str(result))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 27.578551,
   "end_time": "2021-07-07T12:06:31.934681",
   "environment_variables": {},
   "exception": true,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2021-07-07T12:06:04.356130",
   "version": "2.3.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
